package cn.lagou.dw.flume.udfs;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONException;
import com.alibaba.fastjson.JSONObject;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.ListObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;

import java.util.ArrayList;
import java.util.List;

/**
 * new way to write a Hive UDF
 *
 * @author Alex Zhang
 */
public class ParseJsonIds extends GenericUDF {
    PrimitiveObjectInspector arg1;
    PrimitiveObjectInspector arg2;
    ListObjectInspector out;

    @Override
    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
        assert (arguments.length == 2);
        assert (arguments[0].getCategory() == ObjectInspector.Category.PRIMITIVE);
        arg1 = (PrimitiveObjectInspector) arguments[0];
        assert (arg1.getPrimitiveCategory() == PrimitiveObjectInspector.PrimitiveCategory.STRING);
        arg2 = (PrimitiveObjectInspector) arguments[1];
        assert (arg2.getPrimitiveCategory() == PrimitiveObjectInspector.PrimitiveCategory.STRING);
        out = ObjectInspectorFactory.getStandardListObjectInspector(PrimitiveObjectInspectorFactory.writableStringObjectInspector);
        return out;
    }

    @Override
    public Object evaluate(DeferredObject[] arguments) throws HiveException {
        assert (arguments.length == 2);
        Object objectJsonStr = arguments[0].get();
        Object objectArrayKey = arguments[1].get();
        String jsonStr = (String) arg1.getPrimitiveJavaObject(objectJsonStr);
        String arrayKey = (String) arg1.getPrimitiveJavaObject(objectArrayKey);
        try {
            JSONObject jsonObject = JSON.parseObject(jsonStr);
            JSONArray jsonArray = jsonObject.getJSONArray(arrayKey);
            List<String> rs = new ArrayList<>();
            for (Object o : jsonArray) {
                rs.add(o.toString());
            }
            return rs;
        } catch (JSONException e) {
            return null;
        }
    }

    @Override
    public String getDisplayString(String[] children) {
        return null;
    }
}
